package com.ymh.config;

import com.alibaba.csp.sentinel.slots.block.BlockException;
import com.alibaba.fastjson.JSON;
import com.ymh.common.domain.ResultData;
import org.springframework.core.Ordered;
import org.springframework.core.annotation.Order;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.WebExceptionHandler;
import reactor.core.publisher.Mono;

/**
 * @author: 13590
 * @date: 2020/11/10 20:45
 * @description:
 */
@Component
@Order(Ordered.HIGHEST_PRECEDENCE)
public class CustomGatewayBlockExceptionHandler implements WebExceptionHandler {

  private Mono<Void> writeResponse(ServerWebExchange exchange, Throwable ex ) {
    ServerHttpResponse serverHttpResponse = exchange.getResponse();
    serverHttpResponse.getHeaders().add("Content-Type", "application/json;charset=UTF-8");
    ResultData<Object> resultData = ResultData.error( ex.getMessage());
    String resultString = JSON.toJSONString(resultData);
    DataBuffer buffer = serverHttpResponse.bufferFactory().wrap(resultString.getBytes());
    return serverHttpResponse.writeWith(Mono.just(buffer));
  }

  @Override
  public Mono<Void> handle(ServerWebExchange exchange, Throwable ex) {
    if (exchange.getResponse().isCommitted()) {
      return Mono.error(ex);
    } else {
      return !BlockException.isBlockException(ex) ? Mono.error(ex) : this.writeResponse(exchange,ex);
    }
  }
}
